#include "ChromeSession.h"

static const wchar_t g_cszUserDataDir[] = L"--user-data-dir=\"%s\"";

int WINAPI wWinMain(__in HINSTANCE hInstance, __in HINSTANCE hPrevInstance, __in_z LPWSTR lpszCommandLine, __in int nShowCommand)
{
  UNREFERENCED_PARAMETER(hInstance);
  UNREFERENCED_PARAMETER(hPrevInstance);
  UNREFERENCED_PARAMETER(lpszCommandLine);
  UNREFERENCED_PARAMETER(nShowCommand);

  int     nRET = EXIT_FAILURE;
  LPWSTR  lpszPath, lpszRandom, lpszFlag;
  DWORD   cdwSize;

  lpszPath = (LPWSTR)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(wchar_t) * MAX_PATH);
  if (lpszPath != NULL)
  {
    do {
      // get local application data folder
      if (!SHGetSpecialFolderPathW(NULL, lpszPath, CSIDL_LOCAL_APPDATA, FALSE))
        break;

      // set up the directory
      if (!PathAppendW(lpszPath, L"Google\\Chrome"))  break;
      if (!SetCurrentDirectoryW(lpszPath))            break;

      // if CTRL isn't pressed, just run Chrome and exit
      if (!(GetAsyncKeyState(VK_CONTROL) & 0x8000))
      {
        if ((int)ShellExecuteW(NULL, L"open", L"Application\\chrome.exe", NULL, NULL, SW_NORMAL) > 32)
          nRET = EXIT_SUCCESS;
        break;
      }

      // append rest of the path
      if (!PathAppendW(lpszPath, L"User Data"))       break;
      if (!PathAppendW(lpszPath, L"ChromeSession"))   break;

      // is either ALT button pressed? delete everything inside custom user data directory
      if (GetAsyncKeyState(VK_LMENU) & 0x8000 || GetAsyncKeyState(VK_RMENU) & 0x8000)
        DeleteAllInDirectory(lpszPath);

      lpszRandom = (LPWSTR)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(wchar_t) * (ENTROPY + 1));
      if (lpszRandom != NULL)
      {
        do {
          // generate random characters for the new user data directory
          GenerateRandomString(lpszRandom, ENTROPY + 1);
          if (!PathAppendW(lpszPath, lpszRandom))
            break;

          // flag to specify custom user data directory
          cdwSize = lstrlenW(g_cszUserDataDir) - 2 + lstrlenW(lpszPath) + 1;
          lpszFlag = (LPWSTR)HeapAlloc(GetProcessHeap(), 0, sizeof(wchar_t) * cdwSize);
          if (lpszFlag != NULL)
          {
            do {
              if (FAILED(StringCchPrintfW(lpszFlag, cdwSize, g_cszUserDataDir, lpszPath)))
                break;

              // run chrome with flag specifying our directory
              if ((int)ShellExecuteW(NULL, L"open", L"Application\\chrome.exe", lpszFlag, NULL, SW_NORMAL) > 32)
                nRET = EXIT_SUCCESS;
            } while (0);
            HeapFree(GetProcessHeap(), 0, lpszFlag);
          }
        } while (0);
        HeapFree(GetProcessHeap(), 0, lpszRandom);
      }
    } while (0);
    HeapFree(GetProcessHeap(), 0, lpszPath);
  }
  return nRET;
}